from typing import Any, List, Dict, Optional, Union
from kubernetes import client, config
from src.logging import get_logger
from src.config import get_config

logger = get_logger("service_discovery")


class K8sServiceDiscovery:
    def __init__(self, namespace: Optional[str] = None, label_selector: Optional[Union[str, List[str]]] = None):
        """初始化Kubernetes服务发现模块"""
        # 使用统一配置管理器
        config_manager = get_config()
        service_config = config_manager.service_discovery
        
        self.namespace = namespace or service_config.namespace
        # 支持字符串和数组两种格式
        if label_selector is not None:
            if isinstance(label_selector, str):
                self.label_selector = [label_selector]
            else:
                self.label_selector = label_selector
        else:
            self.label_selector = service_config.label_selector
        self.port_name = service_config.port_name
        self.vllm_pods: List[Dict] = []
        
        # 配置Kubernetes客户端
        # 在集群内运行时使用服务账户
        config.load_incluster_config()
        
        self.core_v1 = client.CoreV1Api()
        self.refresh_services()
        
        # 启动事件监听（可选，用于实时更新）
        # self.start_watching()
    
    def refresh_services(self) -> List[Dict]:
        """获取所有匹配标签的vLLM后端Pod"""
        try:
            # 先获取命名空间中的所有Pod（一次API调用）
            try:
                all_pods = self.core_v1.list_namespaced_pod(namespace=self.namespace)
                logger.debug(f"命名空间 {self.namespace} 中共有 {len(all_pods.items)} 个Pod")
            except Exception as e:
                logger.error(f"获取Pod列表失败: {str(e)}", exc_info=True)
                return []
            
            self.vllm_pods = []
            # 遍历所有Pod，检查是否匹配任何一个label selector
            for pod in all_pods.items:
                # 检查Pod是否匹配任何一个label selector（OR逻辑）
                if not self._matches_any_selector(pod, self.label_selector):
                    continue
                
                # 检查Pod状态
                if not self._is_pod_ready(pod):
                    logger.debug(f"跳过未就绪的Pod: {pod.metadata.name}")
                    continue
                
                # 获取Pod IP
                pod_ip = pod.status.pod_ip
                if not pod_ip:
                    logger.debug(f"Pod {pod.metadata.name} 没有IP地址，跳过")
                    continue
                
                # 获取Pod的容器端口信息（仅保留与配置端口名称匹配的端口）
                matched_port = None
                if pod.spec and pod.spec.containers:
                    for container in pod.spec.containers:
                        if not container.ports:
                            continue
                        for port in container.ports:
                            if port.name == self.port_name:
                                matched_port = port.container_port
                                break
                        if matched_port:
                            break
                if not matched_port:
                    logger.debug(
                        "Pod 未找到匹配的端口，将跳过该Pod",
                        pod_name=pod.metadata.name,
                        expected_port=self.port_name
                    )
                    continue
                
                self.vllm_pods.append({
                    "host": pod_ip,
                    "ip": pod_ip,
                    "port": matched_port,
                    "labels": pod.metadata.labels or {},
                    "type": "Pod",
                    "ready": True
                })
            
            logger.info(f"发现 {len(self.vllm_pods)} 个可用的vLLM后端Pod")
            for pod in self.vllm_pods:
                logger.debug(f"Pod: {pod['host']} IP={pod['ip']} {pod}")
                
            return self.vllm_pods
            
        except Exception as e:
            logger.error(f"Pod发现失败: {str(e)}", exc_info=True)
            return []
    
    def refresh_endpoints(self) -> List[Dict]:
        """兼容性方法，保持向后兼容"""
        return self.refresh_services()
        
    @staticmethod
    def _matches_any_selector(pod, selectors: List[str]) -> bool:
        """检查Pod是否匹配任何一个label selector（OR逻辑）"""
        if not selectors:
            return True
        
        pod_labels: Any | dict[Any, Any] = pod.metadata.labels or {}
        
        for selector in selectors:
            if K8sServiceDiscovery._matches_selector(pod_labels, selector):
                return True
        return False
    
    @staticmethod
    def _matches_selector(labels: Dict[str, str], selector: str) -> bool:
        """检查标签是否匹配label selector（仅支持 key=value 格式）"""
        if not selector:
            return False
        
        selector = selector.strip()
        if '=' not in selector:
            logger.warning("无效的 label selector 格式，应为 key=value", selector=selector)
            return False
        
        key, value = selector.split('=', 1)
        key = key.strip()
        value = value.strip()
        
        if not key:
            logger.warning("label selector 缺少 key", selector=selector)
            return False
        
        return labels.get(key) == value
    
    @staticmethod
    def _is_pod_ready(pod) -> bool:
        """检查Pod是否就绪"""
        # 检查Pod的phase是否为Running
        if pod.status.phase != "Running":
            return False
        
        # 检查Pod的容器状态
        if not pod.status.container_statuses:
            return False
        
        # 检查所有容器是否就绪
        for container_status in pod.status.container_statuses:
            if not container_status.ready:
                return False
            # 检查容器是否正在运行
            if not container_status.state or not container_status.state.running:
                return False
        
        return True